import numpy as np
import timeit

# 定义你的Python函数
def my_py_func(x):
    return x * x + 1

# 使用 frompyfunc 创建ufunc
my_ufunc = np.frompyfunc(my_py_func, 1, 1)

# 使用 vectorize 创建相同功能的函数
my_vectorized_func = np.vectorize(my_py_func)

# 直接使用 NumPy 的内置向量化操作
def my_builtin_ufunc(x):
    return np.square(x) + 1

print(type(my_ufunc),type(my_vectorized_func),type(my_builtin_ufunc))
# 随机生成一些数据
data = np.random.rand(100000)

# 测量执行时间
frompyfunc_time = timeit.timeit('my_ufunc(data)', globals=globals(), number=100)
vectorize_time = timeit.timeit('my_vectorized_func(data)', globals=globals(), number=100)
builtin_ufunc_time = timeit.timeit('my_builtin_ufunc(data)', globals=globals(), number=100)

print(f'frompyfunc time: {frompyfunc_time} seconds')
print(f'vectorize time: {vectorize_time} seconds')
print(f'builtin ufunc time: {builtin_ufunc_time} seconds')